package cn.wubo.xianxiacomment

import cn.wubo.xianxiacomment.utils.addCommentBefore
import com.intellij.openapi.actionSystem.ActionUpdateThread
import com.intellij.openapi.actionSystem.AnAction
import com.intellij.openapi.actionSystem.AnActionEvent
import com.intellij.openapi.actionSystem.CommonDataKeys
import com.intellij.openapi.command.WriteCommandAction
import com.intellij.openapi.progress.ProgressIndicator
import com.intellij.openapi.progress.Task
import com.intellij.psi.PsiMethod
import okhttp3.*
import okhttp3.MediaType.Companion.toMediaType
import okhttp3.RequestBody.Companion.toRequestBody
import org.json.JSONObject

class GenerateXianXiaCommentAction : AnAction() {
    private val client = OkHttpClient()

    override fun actionPerformed(e: AnActionEvent) {
        val project = e.project ?: return
        val editor = e.getData(CommonDataKeys.EDITOR) ?: return
        val psiFile = e.getData(CommonDataKeys.PSI_FILE) ?: return
        val offset = editor.caretModel.offset
        val element = psiFile.findElementAt(offset)?.parent

        if (element is PsiMethod) {
            val methodText = element.text
            val settings = OllamaConfigService.instance
            
            object : Task.Backgroundable(project, "正在生成仙侠注释...", true) {
                override fun run(indicator: ProgressIndicator) {
                    indicator.text = "正在调用 Ollama 服务..."
                    val comment = callOllama(settings.ollamaUrl, settings.modelName, methodText)
                    
                    WriteCommandAction.runWriteCommandAction(project) {
                        insertComment(element, comment)
                    }
                }
                
                override fun onThrowable(error: Throwable) {
                    // Handle error
                }
            }.queue()
        }
    }

    override fun update(e: AnActionEvent) {
        val psiElement = e.getData(CommonDataKeys.PSI_ELEMENT)
        e.presentation.isEnabledAndVisible = psiElement is PsiMethod
    }
    
    override fun getActionUpdateThread(): ActionUpdateThread {
        return ActionUpdateThread.BGT
    }

    private fun callOllama(url: String, model: String, methodText: String): String {
        val prompt = "请为以下方法生成仙侠风格的注释，只返回注释内容，不要包含任何其他内容：\n\n$methodText"
        
        val json = JSONObject()
        json.put("model", model)
        json.put("prompt", prompt)
        json.put("stream", false)
        
        val body = json.toString().toRequestBody("application/json; charset=utf-8".toMediaType())
        val request = Request.Builder()
            .url("$url/api/generate")
            .post(body)
            .build()
            
        client.newBuilder()
            .connectTimeout(600, java.util.concurrent.TimeUnit.SECONDS)
            .build()
            .newCall(request).execute().use { response ->
            val responseBody = response.body?.string()
            if (response.isSuccessful && responseBody != null) {
                val jsonResponse = JSONObject(responseBody)
                return jsonResponse.getString("response")
            }
            return "// 未能生成注释"
        }
    }

    private fun insertComment(method: PsiMethod, comment: String) {
        method.addCommentBefore(comment)
    }
}